# include "scene.h"
# include "discr.h"
# include <stdio.h>
# include <math.h>
# include <stdlib.h>
/*
# include <time.h>
# include <sys/timeb.h>*/

# ifdef _OMP_VEL_ADVANCE_
# include <omp.h>
# endif
/*
# ifdef CALL_SOLVE_POISSON_OMP
extern long int solve_poisson_omp(struct grid *g, const double ***d, unsigned short iRK, double ***fspace);
# else*/
extern long int solve_poisson(struct grid *g, const double ***d, unsigned short iRK);
// # endif

extern void vel_bc(struct grid *g, short tflag, unsigned short iRK);
extern void press_bc(const struct grid *g, short tflag, unsigned short iRK);

extern double dens(double T);
extern double kvisc(double T);

static double getA(int i, int j, int k, const struct grid *g);
static double getB(int i, int j, int k, const struct grid *g);
static double getC(int i, int j, int k, const struct grid *g);
static double get_rA_center(int j, int k, const struct grid *g);

struct diverg_max
     {
      short im, jm, km;
      double divmax;
     };

static void divergence(const struct grid *g, double ***diverg, struct diverg_max *mdiv);

extern unsigned int vel_advance(struct grid *g, unsigned int *pcnt_max)
{
 unsigned int pcnt, pcnt_subttl;
 short i, j, k, _i, _j, _k, nr, nth, nz;

 double density0, *r, *th, *z, ***d, ***D0, ***p, ***_u[nRK+1], ***_v[nRK+1], ***_w[nRK+1], ***_p[nRK],\
        ***_A[nRK], ***_B[nRK], ***_C[nRK], ***divH0, ***divH1, divH2, ***div2p0, div2p1, dt, **rA_center[nRK];

// struct timeb t1, t2;

#ifdef TRACE_FILE
 fprintf(g->trptr, "Entering into 'vel_advance()'...\n");
#endif

 r = g->r; th = g->th; z = g->z; dt = g->dt;
 nr = g->nr; nth = g->nth; nz = g->nz;
 *pcnt_max = pcnt_subttl = 0;
 density0 = g->density0;

  _u[0] = g->u;           _v[0] = g->v;            _w[0] = g->w;           _p[0] = g->p;
  _u[1] = g->arr.q[4];    _v[1] = g->arr.q[5];     _w[1] = g->arr.q[6];    _p[1] = g->arr.q[7];
  _u[2] = g->arr.q[8];    _v[2] = g->arr.q[9];     _w[2] = g->arr.q[10];   _p[2] = g->arr.q[11];
  _u[3] = g->arr.q[12];   _v[3] = g->arr.q[13];    _w[3] = g->arr.q[14];

  D0 = g->arr.q[11]; divH0 = g->arr.q[12]; divH1 = g->arr.q[13]; div2p0 = g->arr.q[14];

  for(i = 0; i < nRK; i++)
    {
     _A[i] = g->memFace_r.q[i];
     _B[i] = g->memFace_th.q[i];
     _C[i] = g->memFace_z.q[i];
    }

/* Calculating divergence field at the current n-th time-step. It gets filled in D0[][][]. */
  divergence((const struct grid *)g, D0, (struct diverg_max *)0);

/* --------------------------- Calculation of u(1), v(1), w(1), and p(1) ---------------------------------- */
  rA_center[0] = _A[0][0];
// ftime(&t1);
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(i,j,k) schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz)
             {
              if(i == 0) rA_center[0][j][k] = get_rA_center(NthG+j, NzG+k, (const struct grid *)g);
              else _A[0][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
             }
          if(i < nr && k < nz)  _B[0][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[0][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(_i,_j,_k,i,j,k) schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[1][i][j][k] = (i == NrG+nr-1) ? 0:_u[0][i][j][k]\
              + C2*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2]-r[i]) );

          _v[1][i][j][k] = _v[0][i][j][k]\
             + C2*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) );

          _w[1][i][j][k] = (k == NzG+nz-1) ? 0:_w[0][i][j][k]\
             + C2*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) );
         }

  g->u = _u[1]; g->v = _v[1]; g->w = _w[1];
  vel_bc(g, 0, 1);

/* _u[1], _v[1], _w[1] is used in "press_bc()". Therefore, the memory used by these can't be distroyed. */

  rA_center[1] = _A[1][0];
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(i,j,k) schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz)
               {
                if(i == 0) rA_center[1][j][k] = get_rA_center(NthG+j, NzG+k, (const struct grid *)g);
                else _A[1][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
               }
          if(i < nr && k < nz)  _B[1][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[1][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

/* Calculating RHS of the PPE-1. */
  p = _p[0]; d = g->arr.q[8];
# ifdef _OPENMP
  #pragma omp parallel for private(_i,_j,_k,i,j,k) schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
      {
       div2p0[i][j][k] = DIV2P_ijk;

       divH0[i][j][k] = 2*(( r[i+1]*_A[0][_i+1][_j][_k] - (_i == 0 ? rA_center[0][_j][_k]:r[i]*_A[0][_i][_j][_k]))/(r[i+1]-r[i])\
                           + (_B[0][_i][_j+1][_k] - _B[0][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                           + (_C[0][_i][_j][_k+1] - _C[0][_i][_j][_k])/(z[k+1]-z[k]);

       divH1[i][j][k] = 2*( (r[i+1]*_A[1][_i+1][_j][_k] - (_i == 0 ? rA_center[1][_j][_k]:r[i]*_A[1][_i][_j][_k]))/(r[i+1]-r[i])\
                          + (_B[1][_i][_j+1][_k] - _B[1][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                          + (_C[1][_i][_j][_k+1] - _C[1][_i][_j][_k])/(z[k+1]-z[k]);

       d[i][j][k] = density0*(divH1[i][j][k] + (D0[i][j][k]/dt + a31*(divH0[i][j][k] - div2p0[i][j][k]/density0))/a32);
      }


/* Solving the PPE(1) to calculate p(1). */

/* Copying _p[0] into _p[1] to provide it as initial condition for pressure while solving the PPE. */
   for(i = 0; i < 2*NrG+nr; i++) for(j = 0; j < 2*NthG+nth; j++) for(k = 0; k < 2*NzG+nz; k++)
   _p[1][i][j][k] = _p[0][i][j][k];

   fprintf(g->dptr, "  [ PPE - 1 ]\n");
   fprintf(g->tptr, "PPE-1:-[( "); fflush(g->tptr);

   g->p = _p[1];

/*** "g->arr.q[9]" is used in this function as temporary space if OPENMP calculation. ***/
// #  ifdef CALL_SOLVE_POISSON_OMP
//    pcnt = solve_poisson_omp(g, (const double ***)d, 1, g->arr.q[9]);
// #  else
   pcnt = solve_poisson(g, (const double ***)d, 1/*, g->arr.q[9]*/);
// #  endif

   fprintf(g->tptr, ")]; ");  fflush(g->tptr);

   pcnt_subttl += pcnt;
   if(pcnt > *pcnt_max) *pcnt_max = pcnt;
/* ------------------------------------------------------------------------------------------------------------------------- */


/* ------------------------------------- Calculation of u(2), v(2), w(2), and p(2) ----------------------------------------- */
/* Calculating u(2), v(2), and w(2). */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(_i,_j,_k,i,j,k) schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[2][i][j][k] = (i == NrG+nr-1) ? 0:_u[0][i][j][k]\
                + a31*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2] - r[i]) )\
                + a32*dt*( _A[1][_i+1][_j][_k] - (2.0/density0)*(_p[1][i+1][j][k] - _p[1][i][j][k])/(r[i+2] - r[i]) );

          _v[2][i][j][k] = _v[0][i][j][k]\
                + a31*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/( (r[i]+r[i+1])*(th[j+2]-th[j]) ) )\
                + a32*dt*( _B[1][_i][_j+1][_k] - (4.0/density0)*(_p[1][i][j+1][k] - _p[1][i][j][k])/( (r[i]+r[i+1])*(th[j+2]-th[j]) ) );

          _w[2][i][j][k] = (k == NzG+nz-1) ? 0:_w[0][i][j][k]\
                               + a31*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) )\
                               + a32*dt*( _C[1][_i][_j][_k+1] - (2.0/density0)*(_p[1][i][j][k+1] - _p[1][i][j][k])/(z[k+2]-z[k]) );
         }

  g->u = _u[2]; g->v = _v[2]; g->w = _w[2];
  vel_bc(g, 0, 2);
/* _u[1], _v[1], _w[1] is used in "press_bc()". Therefore, the memory used by these can't be distroyed. */

  rA_center[2] = _A[2][0];
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(i,j,k) schedule(guided)
# endif
  for(k = 0; k <= nz; k++) for(i = 0; i <= nr; i++) for(j = 0; j <= nth; j++)
         {
          if(j < nth && k < nz)
              {
               if(i == 0) rA_center[2][j][k] = get_rA_center(NthG+j, NzG+k, (const struct grid *)g);
               else _A[2][i][j][k] = getA(NrG+i-1, NthG+j, NzG+k, (const struct grid *)g);
              }
          if(i < nr && k < nz)  _B[2][i][j][k] = getB(NrG+i, NthG+j-1, NzG+k, (const struct grid *)g);
          if(i < nr && j < nth) _C[2][i][j][k] = getC(NrG+i, NthG+j, NzG+k-1, (const struct grid *)g);
         }

 /* Calculation of R.H.S. of the PPE(2). */
  p = _p[1]; d = g->arr.q[4];
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(_i,_j,_k,i,j,k,div2p1,divH2) schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
      {
       div2p1 = DIV2P_ijk;
       divH2 = 2*( (r[i+1]*_A[2][_i+1][_j][_k] - (_i == 0 ? rA_center[2][_j][_k]:r[i]*_A[2][_i][_j][_k]))/(r[i+1]-r[i])\
                 + (_B[2][_i][_j+1][_k] - _B[2][_i][_j][_k])/(th[j+1]-th[j]) )/(r[i]+r[i+1])\
                 + (_C[2][_i][_j][_k+1] - _C[2][_i][_j][_k])/(z[k+1]-z[k]);

       d[i][j][k] = density0*( divH2 + (D0[i][j][k]/dt + a41*(divH0[i][j][k] - div2p0[i][j][k]/density0)\
                                                         + a42*(divH1[i][j][k] - div2p1/density0))/a43 );
      }

/* Solving the PPE(2) to calculate p(2). */

/* Copying _p[1] into _p[2] to provide it as initial condition for pressure while solving the PPE. */
   for(i = 0; i < 2*NrG+nr; i++) for(j = 0; j < 2*NthG+nth; j++) for(k = 0; k < 2*NzG+nz; k++)
      _p[2][i][j][k] = _p[1][i][j][k];

   fprintf(g->dptr, "\n  [ PPE - 2 ]\n");
   fprintf(g->tptr, "PPE-2:-[( "); fflush(g->tptr);

   g->p = _p[2];
/*** "g->arr.q[5]" is used in this function as temporary space if OPENMP calculation. ***/
// #  ifdef CALL_SOLVE_POISSON_OMP
//    pcnt = solve_poisson_omp(g, (const double ***)d, 2, g->arr.q[5]);
// #  else
   pcnt = solve_poisson(g, (const double ***)d, 2/*, g->arr.q[5]*/);
// #  endif
   fprintf(g->tptr, ")]; ");  fflush(g->tptr);

   pcnt_subttl += pcnt;
   if(pcnt > *pcnt_max) *pcnt_max = pcnt;
/* ---------------------------------------------------------------------------------------------------------- */

/* ================================= Calculation of u, v, w, and p at next time-step ========================== */
/* Calculation of u, v, and w at next time-step. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel for private(_i,_j,_k,i,j,k) schedule(guided)
# endif
  for(_k = 0; _k < nz; _k++) for(k = _k+NzG, _i = 0; _i < nr; _i++) for(i = _i+NrG, _j = 0, j = NthG; _j < nth; j++, _j++)
         {
          _u[3][i][j][k] = (i == NrG+nr-1) ? 0:_u[0][i][j][k]\
                               + b1*dt*( _A[0][_i+1][_j][_k] - (2.0/density0)*(_p[0][i+1][j][k] - _p[0][i][j][k])/(r[i+2]-r[i]) )\
                               + b2*dt*( _A[1][_i+1][_j][_k] - (2.0/density0)*(_p[1][i+1][j][k] - _p[1][i][j][k])/(r[i+2]-r[i]) )\
                               + b3*dt*( _A[2][_i+1][_j][_k] - (2.0/density0)*(_p[2][i+1][j][k] - _p[2][i][j][k])/(r[i+2]-r[i]) );

          _v[3][i][j][k] = _v[0][i][j][k]\
                  + b1*dt*( _B[0][_i][_j+1][_k] - (4.0/density0)*(_p[0][i][j+1][k] - _p[0][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) )\
                  + b2*dt*( _B[1][_i][_j+1][_k] - (4.0/density0)*(_p[1][i][j+1][k] - _p[1][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) )\
                  + b3*dt*( _B[2][_i][_j+1][_k] - (4.0/density0)*(_p[2][i][j+1][k] - _p[2][i][j][k])/((r[i]+r[i+1])*(th[j+2]-th[j])) );

          _w[3][i][j][k] = (k == NzG+nz-1) ? 0:_w[0][i][j][k]\
                               + b1*dt*( _C[0][_i][_j][_k+1] - (2.0/density0)*(_p[0][i][j][k+1] - _p[0][i][j][k])/(z[k+2]-z[k]) )\
                               + b2*dt*( _C[1][_i][_j][_k+1] - (2.0/density0)*(_p[1][i][j][k+1] - _p[1][i][j][k])/(z[k+2]-z[k]) )\
                               + b3*dt*( _C[2][_i][_j][_k+1] - (2.0/density0)*(_p[2][i][j][k+1] - _p[2][i][j][k])/(z[k+2]-z[k]) );
         }

  g->u = _u[3]; g->v = _v[3]; g->w = _w[3];
  vel_bc(g, 1, 0);

  g->um = g->arr.q[4]; g->vm = g->arr.q[5]; g->wm = g->arr.q[6]; //To store mean velocities [i.e. velocity at time = t+1/2.]

/* Storing mean velocities and copying the newly evaluated velocities to _u[0], _v[0], and _w[0]. */
# ifdef _OMP_VEL_ADVANCE_
  #pragma omp parallel private(i,j,k)
# endif
     {
#     ifdef _OMP_VEL_ADVANCE_
      #pragma omp for schedule(guided) nowait
#     endif
      for(k = 0; k < nz+2*NzG; k++) for(i = 0; i < nr+2*NrG; i++) for(j = 0; j < nth+2*NthG; j++)
         {
       /* Storing mean-velocities. */
          g->um[i][j][k] = 0.5*(_u[0][i][j][k] + _u[3][i][j][k]);
          g->vm[i][j][k] = 0.5*(_v[0][i][j][k] + _v[3][i][j][k]);
          g->wm[i][j][k] = 0.5*(_w[0][i][j][k] + _w[3][i][j][k]);

       /* Copying newly evaluated velocities. */
          _u[0][i][j][k] = _u[3][i][j][k];
          _v[0][i][j][k] = _v[3][i][j][k];
          _w[0][i][j][k] = _w[3][i][j][k];
         }

#     ifdef _OMP_VEL_ADVANCE_
      #pragma omp for schedule(guided)
#     endif
      for(k = NzG; k < nz + NzG; k++) for(i = NrG; i < nr + NrG; i++) for(j = NthG; j < nth + NthG; j++)
       /* Pressure at the next time-step will be given by _p[0]. */
          _p[0][i][j][k] = d1*_p[0][i][j][k] + d2*_p[1][i][j][k] + d3*_p[2][i][j][k];
     }

  g->u = _u[0]; g->v = _v[0]; g->w = _w[0];  g->p = _p[0];
  press_bc(g, 1, 0);
/* ---------------------------------------------------------------------------------------------------------------- */

#ifdef TRACE_FILE
 fprintf(g->trptr, "...'vel_advance()' ends.\n");
#endif
 return pcnt_subttl;

/* q[0]~q[6] are engagged. These can't be used in other functions. */
}



static double getA(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0;

  if(i == NrG-1)
     {
      fprintf(g->lptr, "ERROR:: In 'getA()': i = NrG-1 = %d: The fuction is designed only for i >= NrG.\n", NrG-1);
      exit(-1);
     }

  u = g->u;  v = g->v;  w = g->w; T = g->T; r = g->r;  th = g->th;  z = g->z;
  kvisc0 = g->kvisc0;

  return AA;
 }

static double get_rA_center(int j, int k, const struct grid *g)
 {
  double ***u, ***v, /****w, */ ***T, *r, *th, /* *z, */ kvisc0;
  int i;

  u = g->u;  v = g->v;  /*w = g->w; */T = g->T; r = g->r;  th = g->th;  /*z = g->z;*/
  kvisc0 = g->kvisc0;

  i = NrG-1; //The cell on the face of which "u_center" is defined.

  return rA_CENTER;
 }


static double getB(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0;

  u = g->u;  v = g->v;  w = g->w; T = g->T;  r = g->r;  th = g->th;  z = g->z;
  kvisc0 = g->kvisc0;

  return BB;
 }


static double getC(int i, int j, int k, const struct grid *g)
 {
  double ***u, ***v, ***w, ***T, *r, *th, *z, kvisc0, density0;

  r = g->r;  th = g->th;  z = g->z; u = g->u;  v = g->v;  w = g->w; T = g->T;
  kvisc0 = g->kvisc0; density0 = g->density0;

  return CC;
 }


static void divergence(const struct grid *g, double ***diverg, struct diverg_max *mdiv)
   {
    double *r, *th, *z, ***u, ***v, ***w, divmax, div;
    short i, j, k, im, jm, km, tid;
    struct diverg_max *tddiv;

#   ifdef TRACE_FILE
    fprintf(g->trptr, "Entering into 'divergence()'...\n");
#   endif

    u = g->u;  v = g->v;  w = g->w; r = g->r; th = g->th; z = g->z;

    tddiv = (struct diverg_max *)malloc(sizeof(struct diverg_max)*g->nCPU);

  /* Evaluating divergence: */
    tid = 0;
#   ifdef _OMP_VEL_ADVANCE_
    #pragma omp parallel private(i,j,k,div,tid)
#   endif
         {
#         ifdef _OMP_VEL_ADVANCE_
          tid = omp_get_thread_num();
#         endif

          tddiv[tid].im = -1;
          tddiv[tid].jm = -1;
          tddiv[tid].km = -1;
          tddiv[tid].divmax = 0.0;

#         ifdef _OMP_VEL_ADVANCE_
          #pragma omp for schedule(guided)
#         endif
          for(i = NrG; i < g->nr + NrG; i++) for(j = NthG; j < g->nth + NthG; j++) for(k = NzG; k < g->nz + NzG; k++)
             {
              div = DIV_ijk;

           /* Filtering out the maximum divergence. */
              if((div < 0 ? -div:div) > (tddiv[tid].divmax < 0 ? -tddiv[tid].divmax:tddiv[tid].divmax))
                 {
                  tddiv[tid].divmax = div;
                  tddiv[tid].im = i;
                  tddiv[tid].jm = j;
                  tddiv[tid].km = k;
                 }

              diverg[i][j][k] = div;
             }
         }

 /* Calculating the absolute maximum value of the divergence. */
    divmax = tddiv[0].divmax;
    tid = 0;

    for(i = 1; i < g->nCPU; i++)
        if((tddiv[i].divmax < 0 ? -tddiv[i].divmax:tddiv[i].divmax) > (divmax < 0 ? -divmax:divmax))
           {
            divmax = tddiv[i].divmax;
            tid = i;
           }

    im = tddiv[tid].im;
    jm = tddiv[tid].jm;
    km = tddiv[tid].km;
    free(tddiv);

    fprintf(g->tptr, "divmax(%d, %d, %d) = %E; ", im, jm, km, divmax); fflush(g->tptr);

    if( mdiv != (struct diverg_max *)0 )
         {
          mdiv->divmax = divmax;
          mdiv->im = im;
          mdiv->jm = jm;
          mdiv->km = km;
         }

#   ifdef TRACE_FILE
    fprintf(g->trptr, "...'divergence()' ends.\n");
#   endif

    return;
   }
